import numpy as np
import pandas as pd
from sklearn.feature_selection import VarianceThreshold, SelectKBest, f_classif, chi2, mutual_info_classif
from sklearn.preprocessing import MinMaxScaler
import logging
from scipy.stats import entropy

logger = logging.getLogger(__name__)

class FeatureSelector:
    """特征选择类，实现各种特征选择方法"""
    
    def __init__(self, method='variance', params=None):
        """
        初始化特征选择器
        
        Args:
            method (str): 特征选择方法
                - 'variance': 方差过滤
                - 'correlation': 相关性过滤
                - 'chi2': 卡方检验
                - 'f_test': F检验
                - 'mutual_info': 互信息
                - 'information_gain': 信息增益
                - 'relief': Relief算法
            params (dict): 方法参数
        """
        self.method = method
        self.params = params or {}
        self.selected_features = []
        self.feature_scores = {}
        
    def fit_transform(self, X, y=None):
        """
        执行特征选择并返回选择后的数据
        
        Args:
            X: 特征数据（DataFrame或类数组）
            y: 目标变量（用于监督特征选择方法）
            
        Returns:
            pandas.DataFrame: 选择后的特征数据
        """
        logger.info(f"开始特征选择: 方法={self.method}, 参数={self.params}")
        
        if isinstance(X, pd.DataFrame):
            df = X
        else:
            try:
                df = pd.DataFrame(X)
            except:
                logger.error("输入数据无法转换为DataFrame")
                raise ValueError("输入数据格式无效，无法执行特征选择")
        
        # 记录原始特征
        self.original_features = df.columns.tolist()
        logger.info(f"原始特征数量: {len(self.original_features)}")
        
        try:
            if self.method == 'variance':
                return self._variance_filter(df)
            elif self.method == 'correlation':
                return self._correlation_filter(df)
            elif self.method == 'target_correlation':
                return self._target_correlation_filter(df, self.params.get('target_column'), self.params.get('threshold', 0.1), self.params.get('k'))
            elif self.method == 'chi2':
                return self._chi2_filter(df, y)
            elif self.method == 'f_test':
                return self._f_test_filter(df, y)
            elif self.method == 'mutual_info':
                return self._mutual_info_filter(df, y)
            elif self.method == 'information_gain':
                return self._information_gain_filter(df, y)
            elif self.method == 'relief':
                return self._relief_filter(df, y)
            else:
                logger.error(f"不支持的特征选择方法: {self.method}")
                raise ValueError(f"不支持的特征选择方法: {self.method}")
        except Exception as e:
            logger.exception(f"特征选择过程中出错: {str(e)}")
            raise
    
    def _variance_filter(self, df):
        """方差过滤法"""
        threshold = float(self.params.get('threshold', 0.0))
        logger.info(f"执行方差过滤: 阈值={threshold}")
        
        # 检查是否有数值列
        numeric_columns = df.select_dtypes(include=np.number).columns.tolist()
        if not numeric_columns:
            logger.warning("数据集中没有数值列，无法执行方差过滤")
            self.selected_features = df.columns.tolist()
            return df
        
        # 只对数值列进行方差过滤
        numeric_df = df[numeric_columns]
        non_numeric_df = df.drop(columns=numeric_columns)
        
        # 计算并记录每个特征的方差
        variances = numeric_df.var()
        self.feature_scores = {col: float(var) for col, var in variances.items()}
        
        # 执行方差过滤
        selector = VarianceThreshold(threshold=threshold)
        try:
            selector.fit(numeric_df)
            selected_numeric_columns = numeric_df.columns[selector.get_support()].tolist()
            
            # 合并选中的数值列和非数值列
            self.selected_features = selected_numeric_columns + non_numeric_df.columns.tolist()
            result_df = pd.concat([numeric_df[selected_numeric_columns], non_numeric_df], axis=1)
            
            logger.info(f"方差过滤后保留了 {len(selected_numeric_columns)}/{len(numeric_columns)} 个数值特征")
            return result_df
            
        except Exception as e:
            logger.exception(f"方差过滤失败: {str(e)}")
            # 如果过滤失败，返回原始数据
            self.selected_features = df.columns.tolist()
            return df
    
    def _correlation_filter(self, df):
        """相关性过滤法"""
        threshold = float(self.params.get('threshold', 0.9))
        logger.info(f"执行相关性过滤: 阈值={threshold}")
        
        # 检查是否有数值列
        numeric_columns = df.select_dtypes(include=np.number).columns.tolist()
        if not numeric_columns:
            logger.warning("数据集中没有数值列，无法执行相关性过滤")
            self.selected_features = df.columns.tolist()
            return df
        
        # 只对数值列进行相关性过滤
        numeric_df = df[numeric_columns]
        non_numeric_df = df.drop(columns=numeric_columns)
        
        try:
            # 计算相关系数矩阵
            corr_matrix = numeric_df.corr().abs()
            self.feature_scores = {col: float(corr_matrix[col].mean()) for col in corr_matrix.columns}
            
            # 获取相关系数矩阵的上三角部分
            upper_triangle = np.triu(np.ones(corr_matrix.shape), k=1).astype(bool)
            
            # 找出高相关列
            high_corr_cols = []
            for i, col in enumerate(corr_matrix.columns):
                for j in range(i + 1, len(corr_matrix.columns)):
                    if upper_triangle[i, j] and corr_matrix.iloc[i, j] > threshold:
                        # 如果两列高度相关，选择方差较小的一列删除
                        col_i = corr_matrix.columns[i]
                        col_j = corr_matrix.columns[j]
                        if numeric_df[col_i].var() < numeric_df[col_j].var():
                            high_corr_cols.append(col_i)
                        else:
                            high_corr_cols.append(col_j)
            
            # 删除重复项
            high_corr_cols = list(set(high_corr_cols))
            
            # 保留的特征
            selected_numeric_cols = [col for col in numeric_columns if col not in high_corr_cols]
            
            # 合并选中的数值列和非数值列
            self.selected_features = selected_numeric_cols + non_numeric_df.columns.tolist()
            result_df = pd.concat([numeric_df[selected_numeric_cols], non_numeric_df], axis=1)
            
            logger.info(f"相关性过滤后保留了 {len(selected_numeric_cols)}/{len(numeric_columns)} 个数值特征")
            return result_df
            
        except Exception as e:
            logger.exception(f"相关性过滤失败: {str(e)}")
            # 如果过滤失败，返回原始数据
            self.selected_features = df.columns.tolist()
            return df
    
    def _chi2_filter(self, df, y):
        """卡方检验特征选择"""
        k = int(self.params.get('k', min(10, df.shape[1])))
        logger.info(f"执行卡方检验特征选择: k={k}")
        
        if y is None:
            logger.error("卡方检验需要提供目标变量")
            raise ValueError("卡方检验需要提供目标变量")
        
        # 检查是否有数值列
        numeric_columns = df.select_dtypes(include=np.number).columns.tolist()
        if not numeric_columns:
            logger.warning("数据集中没有数值列，无法执行卡方检验")
            self.selected_features = df.columns.tolist()
            return df
        
        # 只使用数值列进行卡方检验，确保所有值都是非负的
        numeric_df = df[numeric_columns]
        non_numeric_df = df.drop(columns=numeric_columns)
        
        # 对于有负值的列，进行最小-最大缩放到 [0, 1] 区间
        for col in numeric_columns:
            if numeric_df[col].min() < 0:
                logger.info(f"列 '{col}' 包含负值，进行最小-最大缩放")
                numeric_df[col] = (numeric_df[col] - numeric_df[col].min()) / (numeric_df[col].max() - numeric_df[col].min())
        
        try:
            # 限制k不超过特征数量
            k = min(k, len(numeric_columns))
            
            # 执行卡方检验
            selector = SelectKBest(score_func=chi2, k=k)
            selector.fit(numeric_df, y)
            
            # 获取特征得分
            scores = selector.scores_
            self.feature_scores = {col: float(score) for col, score in zip(numeric_columns, scores)}
            
            # 获取选中的特征
            selected_numeric_cols = numeric_df.columns[selector.get_support()].tolist()
            
            # 合并选中的数值列和非数值列
            self.selected_features = selected_numeric_cols + non_numeric_df.columns.tolist()
            result_df = pd.concat([numeric_df[selected_numeric_cols], non_numeric_df], axis=1)
            
            logger.info(f"卡方检验后保留了 {len(selected_numeric_cols)}/{len(numeric_columns)} 个数值特征")
            return result_df
            
        except Exception as e:
            logger.exception(f"卡方检验特征选择失败: {str(e)}")
            # 如果选择失败，返回原始数据
            self.selected_features = df.columns.tolist()
            return df
    
    def _f_test_filter(self, df, y):
        """F检验特征选择"""
        k = int(self.params.get('k', min(10, df.shape[1])))
        logger.info(f"执行F检验特征选择: k={k}")
        
        if y is None:
            logger.error("F检验需要提供目标变量")
            raise ValueError("F检验需要提供目标变量")
        
        # 检查是否有数值列
        numeric_columns = df.select_dtypes(include=np.number).columns.tolist()
        if not numeric_columns:
            logger.warning("数据集中没有数值列，无法执行F检验")
            self.selected_features = df.columns.tolist()
            return df
        
        # 只使用数值列进行F检验
        numeric_df = df[numeric_columns]
        non_numeric_df = df.drop(columns=numeric_columns)
        
        try:
            # 限制k不超过特征数量
            k = min(k, len(numeric_columns))
            
            # 执行F检验
            selector = SelectKBest(score_func=f_classif, k=k)
            selector.fit(numeric_df, y)
            
            # 获取特征得分
            scores = selector.scores_
            self.feature_scores = {col: float(score) for col, score in zip(numeric_columns, scores)}
            
            # 获取选中的特征
            selected_numeric_cols = numeric_df.columns[selector.get_support()].tolist()
            
            # 合并选中的数值列和非数值列
            self.selected_features = selected_numeric_cols + non_numeric_df.columns.tolist()
            result_df = pd.concat([numeric_df[selected_numeric_cols], non_numeric_df], axis=1)
            
            logger.info(f"F检验后保留了 {len(selected_numeric_cols)}/{len(numeric_columns)} 个数值特征")
            return result_df
            
        except Exception as e:
            logger.exception(f"F检验特征选择失败: {str(e)}")
            # 如果选择失败，返回原始数据
            self.selected_features = df.columns.tolist()
            return df
    
    def _mutual_info_filter(self, df, y):
        """互信息特征选择"""
        k = int(self.params.get('k', min(10, df.shape[1])))
        logger.info(f"执行互信息特征选择: k={k}")
        
        if y is None:
            logger.error("互信息特征选择需要提供目标变量")
            raise ValueError("互信息特征选择需要提供目标变量")
        
        # 检查是否有数值列
        numeric_columns = df.select_dtypes(include=np.number).columns.tolist()
        if not numeric_columns:
            logger.warning("数据集中没有数值列，无法执行互信息特征选择")
            self.selected_features = df.columns.tolist()
            return df
        
        # 只使用数值列进行互信息特征选择
        numeric_df = df[numeric_columns]
        non_numeric_df = df.drop(columns=numeric_columns)
        
        try:
            # 限制k不超过特征数量
            k = min(k, len(numeric_columns))
            
            # 执行互信息特征选择
            selector = SelectKBest(score_func=mutual_info_classif, k=k)
            selector.fit(numeric_df, y)
            
            # 获取特征得分
            scores = selector.scores_
            self.feature_scores = {col: float(score) for col, score in zip(numeric_columns, scores)}
            
            # 获取选中的特征
            selected_numeric_cols = numeric_df.columns[selector.get_support()].tolist()
            
            # 合并选中的数值列和非数值列
            self.selected_features = selected_numeric_cols + non_numeric_df.columns.tolist()
            result_df = pd.concat([numeric_df[selected_numeric_cols], non_numeric_df], axis=1)
            
            logger.info(f"互信息特征选择后保留了 {len(selected_numeric_cols)}/{len(numeric_columns)} 个数值特征")
            return result_df
            
        except Exception as e:
            logger.exception(f"互信息特征选择失败: {str(e)}")
            # 如果选择失败，返回原始数据
            self.selected_features = df.columns.tolist()
            return df

    def _target_correlation_filter(self, df, target_column, threshold=0.1, k=None):
        """
        基于与目标变量的相关性进行特征选择
        
        Args:
            df (pd.DataFrame): 输入数据
            target_column (str): 目标列名
            threshold (float): 相关性阈值，低于此阈值的特征将被移除
            k (int): 如果提供，选择相关性最高的k个特征
            
        Returns:
            pd.DataFrame: 过滤后的数据
            list: 选择的特征列表
            dict: 特征得分（相关系数）
        """
        try:
            logger.info(f"执行基于目标相关性的特征选择: 目标列={target_column}")
            
            # 检查目标列是否存在
            if target_column not in df.columns:
                logger.error(f"目标列 {target_column} 不存在")
                return df, df.columns.tolist(), {}
            
            # 从数据中提取目标列
            y = df[target_column]
            
            # 仅保留数值型特征列
            numeric_cols = df.select_dtypes(include=['number']).columns
            X = df[numeric_cols].drop(columns=[target_column]) if target_column in numeric_cols else df[numeric_cols]
            
            if X.empty:
                logger.warning("没有数值型特征可以计算相关性")
                return df, df.columns.tolist(), {}
            
            # 计算与目标列的相关性
            correlation_with_target = {}
            for col in X.columns:
                if col != target_column:
                    corr = X[col].corr(y)
                    if not pd.isna(corr):  # 排除NaN相关系数
                        correlation_with_target[col] = abs(corr)  # 使用绝对值，因为我们关心相关性的强度
            
            logger.info(f"计算了 {len(correlation_with_target)} 个特征与目标的相关性")
            
            # 按相关性排序
            sorted_correlations = sorted(correlation_with_target.items(), key=lambda x: x[1], reverse=True)
            
            # 选择特征
            selected_features = []
            if k is not None and k > 0:
                # 如果指定了k，选择相关性最高的k个特征
                selected_features = [col for col, _ in sorted_correlations[:k]]
                logger.info(f"选择了相关性最高的 {k} 个特征")
            else:
                # 否则，使用阈值过滤
                selected_features = [col for col, corr in sorted_correlations if corr >= threshold]
                logger.info(f"选择了相关性大于等于 {threshold} 的 {len(selected_features)} 个特征")
            
            # 选择列
            final_features = selected_features + [target_column]
            # 创建特征得分字典
            feature_scores = {col: corr for col, corr in correlation_with_target.items()}
            
            return df[final_features], final_features, feature_scores
            
        except Exception as e:
            logger.exception(f"目标相关性过滤失败: {str(e)}")
            return df, df.columns.tolist(), {}

    def _information_gain_filter(self, df, y, k=None):
        """
        基于信息增益的特征选择
        
        Args:
            df (pd.DataFrame): 输入数据
            y (pd.Series): 目标变量
            k (int, 可选): 选择的特征数量
            
        Returns:
            pd.DataFrame: 过滤后的数据
            list: 选择的特征列表
            dict: 特征得分（信息增益）
        """
        try:
            logger.info("执行信息增益特征选择")
            
            if k is None:
                k = int(self.params.get('k', 10))
                
            # 提取目标列名
            target_column = y.name
            
            # 仅保留数值型特征列
            numeric_cols = df.select_dtypes(include=['number']).columns
            
            # 从数值列中排除目标列
            if target_column in numeric_cols:
                numeric_cols = numeric_cols.drop([target_column])
                
            if len(numeric_cols) == 0:
                logger.warning("没有数值型特征可以进行信息增益计算")
                return df, df.columns.tolist(), {}
                
            # 计算目标变量的熵
            def calc_entropy(y_values):
                # 目标变量可能是分类的或连续的，先进行处理
                if len(set(y_values)) <= 10:  # 认为是分类变量
                    # 计算每个类别的概率
                    value_counts = pd.Series(y_values).value_counts(normalize=True)
                    return entropy(value_counts)
                else:  # 认为是连续变量
                    # 对连续变量进行分箱处理
                    bins = min(10, len(set(y_values)))
                    y_binned = pd.cut(y_values, bins=bins)
                    # 计算每个箱子的概率
                    value_counts = y_binned.value_counts(normalize=True)
                    return entropy(value_counts)
                    
            # 计算条件熵
            def calc_conditional_entropy(x_values, y_values):
                # x和y都可能是分类或连续的
                # 处理x
                if len(set(x_values)) <= 10:  # 认为是分类变量
                    # 对每个x的类别计算条件熵
                    cond_entropy = 0
                    for x_val in set(x_values):
                        mask = x_values == x_val
                        prop = sum(mask) / len(x_values)
                        cond_entropy += prop * calc_entropy(y_values[mask])
                    return cond_entropy
                else:  # 认为是连续变量
                    # 对x进行分箱处理
                    bins = min(10, len(set(x_values)))
                    x_binned = pd.cut(x_values, bins=bins)
                    # 对每个箱子计算条件熵
                    cond_entropy = 0
                    for x_bin in set(x_binned):
                        mask = x_binned == x_bin
                        if sum(mask) > 0:  # 避免空箱子
                            prop = sum(mask) / len(x_values)
                            cond_entropy += prop * calc_entropy(y_values[mask])
                    return cond_entropy
            
            # 计算信息增益
            target_entropy = calc_entropy(y.values)
            feature_scores = {}
            
            for col in numeric_cols:
                x_values = df[col].values
                # 避免NaN值
                mask = ~np.isnan(x_values) & ~np.isnan(y.values)
                if sum(mask) > 0:  # 确保有有效值
                    cond_entropy = calc_conditional_entropy(x_values[mask], y.values[mask])
                    info_gain = target_entropy - cond_entropy
                    feature_scores[col] = info_gain
            
            logger.info(f"计算了 {len(feature_scores)} 个特征的信息增益")
            
            # 按信息增益排序
            sorted_features = sorted(feature_scores.items(), key=lambda x: x[1], reverse=True)
            
            # 选择前k个特征
            k = min(k, len(sorted_features))
            selected_features = [col for col, _ in sorted_features[:k]]
            
            logger.info(f"选择了 {len(selected_features)} 个特征")
            
            # 选择列 - 包括原始的非数值列和目标列
            all_columns = df.columns.tolist()
            non_numeric_cols = [col for col in all_columns if col not in numeric_cols and col != target_column]
            final_cols = selected_features + non_numeric_cols + [target_column]
            
            return df[final_cols], final_cols, feature_scores
            
        except Exception as e:
            logger.exception(f"信息增益特征选择失败: {str(e)}")
            return df, df.columns.tolist(), {}
            
    def _relief_filter(self, df, y, k=None, m=None, distance_threshold=None):
        """
        基于Relief算法的特征选择
        
        Args:
            df (pd.DataFrame): 输入数据
            y (pd.Series): 目标变量
            k (int, 可选): 选择的特征数量
            m (int, 可选): 采样的实例数量
            distance_threshold (float, 可选): 距离阈值
            
        Returns:
            pd.DataFrame: 过滤后的数据
            list: 选择的特征列表
            dict: 特征得分（Relief评分）
        """
        try:
            logger.info("执行Relief算法特征选择")
            
            if k is None:
                k = int(self.params.get('k', 10))
                
            if m is None:
                # 默认使用样本数量的10%，但不少于10个，不多于100个
                m = min(max(10, int(len(df) * 0.1)), 100)
                
            # 提取目标列名
            target_column = y.name
            
            # 仅保留数值型特征列
            numeric_cols = df.select_dtypes(include=['number']).columns
            
            # 从数值列中排除目标列
            if target_column in numeric_cols:
                numeric_cols = numeric_cols.drop([target_column])
                
            if len(numeric_cols) == 0:
                logger.warning("没有数值型特征可以进行Relief评分计算")
                return df, df.columns.tolist(), {}
                
            # 提取特征和目标
            X = df[numeric_cols].values
            y = y.values
            
            # 归一化特征，以便计算距离
            scaler = MinMaxScaler()
            X_scaled = scaler.fit_transform(X)
            
            # 检查目标是否是分类问题
            is_classification = len(np.unique(y)) <= 10
            
            # Relief算法
            n_samples, n_features = X_scaled.shape
            weights = np.zeros(n_features)
            
            # 随机选择m个样本
            m = min(m, n_samples)
            indices = np.random.choice(n_samples, m, replace=False)
            
            for i in indices:
                # 当前实例
                instance = X_scaled[i]
                instance_label = y[i]
                
                # 计算与所有其他实例的距离
                distances = np.sqrt(((X_scaled - instance) ** 2).sum(axis=1))
                
                # 根据问题类型找到最近的命中（相同类别）和失误（不同类别）
                if is_classification:
                    # 同类样本的索引
                    same_class = np.where((y == instance_label) & (np.arange(n_samples) != i))[0]
                    # 不同类样本的索引
                    diff_class = np.where(y != instance_label)[0]
                    
                    if len(same_class) > 0 and len(diff_class) > 0:
                        # 找到最近的同类样本（命中）
                        nearest_hit_idx = same_class[np.argmin(distances[same_class])]
                        # 找到最近的不同类样本（失误）
                        nearest_miss_idx = diff_class[np.argmin(distances[diff_class])]
                        
                        # 计算特征对距离的贡献
                        diff_hit = np.abs(X_scaled[nearest_hit_idx] - instance)
                        diff_miss = np.abs(X_scaled[nearest_miss_idx] - instance)
                        
                        # 更新权重
                        weights = weights - diff_hit + diff_miss
                else:
                    # 对于回归问题，使用不同的方法
                    # 找到最近的邻居
                    nearest_idx = np.argsort(distances)[1:6]  # 排除自身，选择最近的5个
                    
                    # 计算目标差异的归一化
                    y_diffs = np.abs(y[nearest_idx] - instance_label)
                    y_diffs = y_diffs / np.max(y_diffs) if np.max(y_diffs) > 0 else y_diffs
                    
                    # 计算特征对距离的贡献，加权
                    for j, idx in enumerate(nearest_idx):
                        diff = np.abs(X_scaled[idx] - instance)
                        weights = weights + diff * y_diffs[j]
            
            # 归一化权重
            if np.max(np.abs(weights)) > 0:
                weights = weights / np.max(np.abs(weights))
                
            # 创建特征分数字典
            feature_scores = {col: weights[i] for i, col in enumerate(numeric_cols)}
            
            # 按权重排序
            sorted_features = sorted(feature_scores.items(), key=lambda x: x[1], reverse=True)
            
            # 选择前k个特征
            k = min(k, len(sorted_features))
            selected_features = [col for col, _ in sorted_features[:k]]
            
            logger.info(f"Relief算法选择了 {len(selected_features)} 个特征")
            
            # 选择列 - 包括原始的非数值列和目标列
            all_columns = df.columns.tolist()
            non_numeric_cols = [col for col in all_columns if col not in numeric_cols and col != target_column]
            final_cols = selected_features + non_numeric_cols + [target_column]
            
            return df[final_cols], final_cols, feature_scores
            
        except Exception as e:
            logger.exception(f"Relief算法特征选择失败: {str(e)}")
            return df, df.columns.tolist(), {}

def feature_selection_core(df, method, params, target_column=None):
    """
    核心特征选择函数
    
    Args:
        df (pd.DataFrame): 输入数据
        method (str): 特征选择方法
        params (dict): 方法参数
        target_column (str, optional): 目标变量列名
    
    Returns:
        tuple: (选择后的数据, 选择的特征列表, 特征得分字典)
    """
    logger.info(f"开始执行特征选择: 方法={method}, 目标列={target_column}")
    
    # 处理目标变量
    y = None
    if target_column and target_column in df.columns:
        y = df[target_column].values
        # 从特征中移除目标变量
        X = df.drop(columns=[target_column])
    else:
        X = df
    
    # 创建特征选择器
    selector = FeatureSelector(method=method, params=params)
    
    # 执行特征选择
    result_df = selector.fit_transform(X, y)
    
    # 如果存在目标变量，将其添加回结果
    if target_column and target_column in df.columns:
        result_df[target_column] = df[target_column]
        # 确保目标列也在选择的特征列表中
        if target_column not in selector.selected_features:
            selector.selected_features.append(target_column)
    
    logger.info(f"特征选择完成: 从 {df.shape[1]} 个特征选择了 {len(selector.selected_features)} 个特征")
    
    # 返回结果DataFrame, 选择的特征列表和特征得分
    return result_df, selector.selected_features, selector.feature_scores 